import os
import sys

import shutil
import subprocess

os.environ["nnUNet_raw"] = "nnUNetFrame/DATASET/nnUNet_raw"
os.environ["nnUNet_preprocessed"] = "nnUNetFrame/DATASET/nnUNet_preprocessed"
os.environ["nnUNet_results"] = "nnUNetFrame/DATASET/nnUNet_results"


def _fix_name(path: str):
    for each_name in os.listdir(path):
        if each_name.endswith(".nii.gz"):
            # hippocampus_001.nii.gz -> hippocampus_001_0000.nii.gz
            ori_name = each_name.split(".nii.gz")[0]
            new_name = "{}_0000.nii.gz".format(ori_name)
            shutil.move(os.path.join(path, each_name), os.path.join(path, new_name))


def fix_dataset_name():
    # 修复数据集名称到规定格式
    # BRATS_001_0000.nii.gz: 任意_第几个_第几个的第几个通道.nii.gz
    dataset_path = "nnUNetFrame/DATASET/nnUNet_raw/Dataset004_Hippocampus"
    imagesTr = os.path.join(dataset_path, 'imagesTr')
    imagesTs = os.path.join(dataset_path, 'imagesTs')

    _fix_name(imagesTr)
    _fix_name(imagesTs)


def preprocess():
    # nnUNet/nnunetv2/experiment_planning/里面的几个py文件
    # nnUNetFrame/DATASET/nnUNet_preprocessed/Dataset004_Hippocampus, 这里会有处理结果生成
    subprocess.check_call(["Depend/Scripts/nnUNetv2_plan_and_preprocess", "-d", "4", "--verify_dataset_integrity"])


def train():
    # nnUNet/nnunetv2/run/run_training.py
    # nnUNet/documentation/how_to_use_nnunet.md
    #   - nnUNetv2_train DATASET_NAME_OR_ID 2d FOLD [--npz]
    #   - nnUNetv2_train DATASET_NAME_OR_ID 3d_fullres FOLD [--npz]
    # nnUNet/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py num_epochs调整轮数
    # 这里为了调试，设置为了1
    subprocess.check_call(["Depend/Scripts/nnUNetv2_train", "4", "3d_fullres", "0", "--npz", '-device', 'cpu'])



if __name__ == '__main__':
    # fix_dataset_name()
    # preprocess()
    train()
